import numpy as np

from core.config import Config
from utils.common import is_variable

try:
    import cupy as cp
    using_gpu = True
    array_types = (np.ndarray, cp.ndarray)
except ImportError:
    using_gpu = False
    array_types = np.ndarray
    print("Warning: Cupy module not found. GPU acceleration is disabled.")


class Cuda:

    # 安装cupy
    # conda install -c conda-forge cupy
    # 确保安装成功cuda和visual studio后，要配置环境
    # CUDA Toolkit的bin目录，例如C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\vxx.x\bin
    # Visual Studio的编译器目录，例如C:\Program Files (x86)\Microsoft Visual Studio\2019\Community\VC\Tools\MSVC\14.xx.xxxxxx\bin\Hostx64\x64

    def __init__(self, gpu_enable=False):
        self.module = cp if using_gpu and gpu_enable else np
        self.gpu_enable = using_gpu and gpu_enable
        self.array_types = array_types

    def get_array_module(self, x):
        return self.module

    def to_array(self, x):
        if self.gpu_enable:
            return self.to_cupy(x)
        else:
            return self.to_numpy(x)


    def to_numpy(self, x):
        """Convert to `numpy.ndarray`.

        Args:
            x (`numpy.ndarray` or `cupy.ndarray`): Arbitrary object that can be
                converted to `numpy.ndarray`.
        Returns:
            `numpy.ndarray`: Converted array.
        """

        if is_variable(x):
            x = x.data

        if np.isscalar(x) or isinstance(x, (list, tuple, int, float)):
            return np.array(x)
        elif isinstance(x, np.ndarray):
            return x
        # cupy -> numpy
        return x.get()

    def to_cupy(self, x):
        """Convert to `cupy.ndarray`.

        Args:
            x (`numpy.ndarray` or `cupy.ndarray`): Arbitrary object that can be
                converted to `cupy.ndarray`.
        Returns:
            `cupy.ndarray`: Converted array.
        """
        if is_variable(x):
            x = x.data

        if not self.gpu_enable:
            raise Exception('CuPy cannot be loaded. Install CuPy!')
        # numpy -> cupy
        return self.module.asarray(x)


# cuda实例
cuda = Cuda(gpu_enable=Config.enable_gpu)
# 引用的模组 np或者cp
cuda_module = cuda.module
cuda_array_types = cuda.array_types
